/*
 * SPDX-FileCopyrightText: 2022-2024 Andrew Gunnerson
 * SPDX-License-Identifier: GPL-3.0-only
 */

@file:OptIn(ExperimentalUnsignedTypes::class)

package com.sanji.audio_test.format

import android.media.MediaCodec
import android.media.MediaFormat
import android.system.Os
import android.system.OsConstants
import com.sanji.audio_test.writeFully
import java.io.FileDescriptor
import java.nio.ByteBuffer
import java.nio.ByteOrder

/**
 * WAVE容器实现类，用于处理WAV音频文件格式
 *
 * @param fd 文件描述符
 */
class WaveContainer(private val fd: FileDescriptor) : Container {
    // 标记容器是否已启动
    private var isStarted = false
    // 轨道索引，初始值为-1表示尚未添加轨道
    private var track = -1
    // 帧大小（每帧字节数）
    private var frameSize = 0
    // 声道数
    private var channelCount = 0
    // 采样率
    private var sampleRate = 0

    /**
     * 启动容器
     * @throws IllegalStateException 如果容器已经启动
     */
    override fun start() {
        // 检查容器是否已经启动，如果已启动则抛出异常
        if (isStarted) {
            throw IllegalStateException("容器已经启动")
        }

        // 截断文件到0字节，清空内容
        Os.ftruncate(fd, 0)

        // 跳过头部位置，为后续写入音频数据做准备
        Os.lseek(fd, HEADER_SIZE.toLong(), OsConstants.SEEK_SET)

        // 设置容器状态为已启动
        isStarted = true
    }

    /**
     * 停止容器
     * @throws IllegalStateException 如果容器尚未启动
     */
    override fun stop() {
        // 检查容器是否未启动，如果未启动则抛出异常
        if (!isStarted) {
            throw IllegalStateException("容器尚未启动")
        }

        // 设置容器状态为未启动
        isStarted = false

        // 如果已添加轨道，则构建并写入WAV文件头
        if (track >= 0) {
            // 获取当前文件大小
            val fileSize = Os.lseek(fd, 0, OsConstants.SEEK_CUR)
            // 构建文件头
            val header = buildHeader(fileSize)
            // 回到文件开始位置
            Os.lseek(fd, 0, OsConstants.SEEK_SET)
            // 写入文件头
            writeFully(fd, header)
        }
    }

    /**
     * 释放容器资源
     * 如果容器正在运行，则先停止它
     */
    override fun release() {
        // 如果容器正在运行，则停止它
        if (isStarted) {
            stop()
        }
    }

    /**
     * 添加音轨
     *
     * @param mediaFormat 媒体格式信息
     * @return 轨道索引
     * @throws IllegalStateException 如果容器已启动或轨道已添加
     */
    override fun addTrack(mediaFormat: MediaFormat): Int {
        // 检查容器是否已启动或轨道是否已添加
        if (isStarted) {
            throw IllegalStateException("容器已经启动")
        } else if (track >= 0) {
            throw IllegalStateException("轨道已经添加")
        }

        // 设置轨道索引
        track = 0
        // 从媒体格式中获取帧大小
        frameSize = mediaFormat.getInteger(Format.KEY_X_FRAME_SIZE_IN_BYTES)
        // 从媒体格式中获取声道数
        channelCount = mediaFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT)
        // 从媒体格式中获取采样率
        sampleRate = mediaFormat.getInteger(MediaFormat.KEY_SAMPLE_RATE)

        // 返回轨道索引
        return track
    }

    /**
     * 写入音频样本数据
     *
     * @param trackIndex 轨道索引
     * @param byteBuffer 包含音频数据的字节缓冲区
     * @param bufferInfo 缓冲区信息
     * @throws IllegalStateException 如果容器未启动、未添加轨道或轨道索引无效
     */
    override fun writeSamples(trackIndex: Int, byteBuffer: ByteBuffer,
                              bufferInfo: MediaCodec.BufferInfo) {
        // 检查容器状态和轨道信息
        if (!isStarted) {
            throw IllegalStateException("容器尚未启动")
        } else if (track < 0) {
            throw IllegalStateException("尚未添加轨道")
        } else if (track != trackIndex) {
            throw IllegalStateException("无效的轨道: $trackIndex")
        }

        // 将音频数据写入文件
        writeFully(fd, byteBuffer)
    }

    /**
     * 构建WAV文件头
     *
     * @param fileSize 文件大小
     * @return 包含WAV文件头的字节缓冲区
     */
    private fun buildHeader(fileSize: Long): ByteBuffer =
        ByteBuffer.allocate(HEADER_SIZE).apply {
            // 设置字节序为小端序
            order(ByteOrder.LITTLE_ENDIAN)

            // 计算块大小和数据大小
            val (chunkSize, dataSize) = if (fileSize >= Int.MAX_VALUE) {
                // 如果文件过大，不设置大小，让音频播放器自行猜测
                Pair(0, 0)
            } else {
                Pair(fileSize.toInt() - 8, fileSize.toInt() - HEADER_SIZE)
            }

            // 0-3: 块ID (RIFF)
            put(RIFF_MAGIC.asByteArray())
            // 4-7: 块大小
            putInt(chunkSize)
            // 8-11: 格式 (WAVE)
            put(WAVE_MAGIC.asByteArray())
            // 12-15: 子块1 ID (fmt )
            put(FMT_MAGIC.asByteArray())
            // 16-19: 子块1大小
            putInt(16)
            // 20-21: 音频格式 (1 = PCM)
            putShort(1)
            // 22-23: 声道数
            putShort(channelCount.toShort())
            // 24-27: 采样率
            putInt(sampleRate)
            // 28-31: 字节率 (采样率 * 帧大小)
            putInt(sampleRate * frameSize)
            // 32-33: 块对齐 (帧大小)
            putShort(frameSize.toShort())
            // 34-35: 每个样本的位数
            putShort(((frameSize / channelCount) * 8).toShort())
            // 36-39: 子块2 ID (data)
            put(DATA_MAGIC.asByteArray())
            // 40-43: 子块2大小
            putInt(dataSize)

            // 翻转缓冲区，准备读取
            flip()
        }

    companion object {
        // WAV文件头大小（44字节）
        private const val HEADER_SIZE = 44
        // RIFF标识符
        private val RIFF_MAGIC = ubyteArrayOf(0x52u, 0x49u, 0x46u, 0x46u) // RIFF
        // WAVE标识符
        private val WAVE_MAGIC = ubyteArrayOf(0x57u, 0x41u, 0x56u, 0x45u) // WAVE
        // fmt 标识符
        private val FMT_MAGIC = ubyteArrayOf(0x66u, 0x6du, 0x74u, 0x20u) // "fmt "
        // data标识符
        private val DATA_MAGIC = ubyteArrayOf(0x64u, 0x61u, 0x74u, 0x61u) // data
    }
}